import os
from typing import List, Dict, Any, Optional

import dashscope
from dashscope.audio.asr import VocabularyService
from dotenv import load_dotenv

from config.logging_config import get_logger

# 加载环境变量
load_dotenv()

# 配置日志
logger = get_logger(__name__)

# 配置DashScope API Key
if not dashscope.api_key:
    dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")

# 音频的热词设定
MY_VOCABULARY = [
    {"text": "路单", "weight": 4, "lang": "zh"},
    {"text": "排班", "weight": 4, "lang": "zh"},
    {"text": "清屏", "weight": 4, "lang": "zh"},
    {"text": "关闭麦克风", "weight": 4, "lang": "zh"},
    {"text": "清除", "weight": 4, "lang": "zh"},
    {"text": "清空", "weight": 4, "lang": "zh"},
    {"text": "重录", "weight": 4, "lang": "zh"},
]

VOCABULARY_ID = "vocab-asr1-c8d82369dd714eba877b4ff10df5aa72"  # 默认热词表ID


class VocabularyManager:
    """热词管理器 - 支持增删改查操作"""

    def __init__(self, api_key: Optional[str] = None):
        """
        初始化热词管理器
        
        Args:
            api_key: DashScope API Key，如果不提供将从环境变量获取
        """
        if api_key:
            dashscope.api_key = api_key
        elif not dashscope.api_key:
            raise ValueError("DashScope API Key 未配置，请设置环境变量 DASHSCOPE_API_KEY 或传入 api_key 参数")

        self.service = VocabularyService()
        self.vocabulary_cache: Dict[str, List[dict]] = {}

    def create_vocabulary(self, target_model: str, prefix: str, vocabulary: List[dict]) -> str:
        """
        创建热词表
        
        Args:
            target_model: 热词表对应的语音识别模型版本 (如: paraformer-realtime-v2)
            prefix: 热词表自定义前缀，仅允许数字和小写字母，小于十个字符
            vocabulary: 热词表字典列表
            
        Returns:
            热词表标识符 vocabulary_id
            
        Example:
            vocabulary = [
                {"text": "吴贻弓", "weight": 4, "lang": "zh"},
                {"text": "阙里人家", "weight": 4, "lang": "zh"}
            ]
        """
        try:
            # 验证参数
            self._validate_vocabulary(vocabulary)
            self._validate_prefix(prefix)

            # 创建热词表
            vocabulary_id = self.service.create_vocabulary(
                prefix=prefix,
                target_model=target_model,
                vocabulary=vocabulary
            )

            # 缓存热词表
            self.vocabulary_cache[vocabulary_id] = vocabulary.copy()

            print(f"热词表创建成功，ID: {vocabulary_id}")
            return vocabulary_id

        except Exception as e:
            print(f"创建热词表失败: {e}")
            raise

    def update_vocabulary(self, vocabulary_id: str, vocabulary: List[dict]) -> bool:
        """
        更新热词表
        
        Args:
            vocabulary_id: 需要更新的热词表ID
            vocabulary: 新的热词表字典列表
            
        Returns:
            是否更新成功
        """
        try:
            # 验证参数
            self._validate_vocabulary(vocabulary)

            # 更新热词表
            self.service.update_vocabulary(vocabulary_id, vocabulary)

            # 更新缓存
            self.vocabulary_cache[vocabulary_id] = vocabulary.copy()

            print(f"热词表更新成功，ID: {vocabulary_id}")
            return True

        except Exception as e:
            print(f"更新热词表失败: {e}")
            return False

    def delete_vocabulary(self, vocabulary_id: str) -> bool:
        """
        删除热词表
        
        Args:
            vocabulary_id: 要删除的热词表ID
            
        Returns:
            是否删除成功
        """
        try:
            # 删除热词表
            self.service.delete_vocabulary(vocabulary_id)

            # 清除缓存
            if vocabulary_id in self.vocabulary_cache:
                del self.vocabulary_cache[vocabulary_id]

            print(f"热词表删除成功，ID: {vocabulary_id}")
            return True

        except Exception as e:
            print(f"删除热词表失败: {e}")
            return False

    def get_vocabulary(self, vocabulary_id: str) -> Optional[Dict[str, Any]]:
        """
        获取单个热词表详情
        
        Args:
            vocabulary_id: 热词表ID
            
        Returns:
            热词表详情字典，如果不存在返回None
        """
        vocabulary_data = []
        try:
            # 优先从缓存获取
            if vocabulary_id in self.vocabulary_cache:
                vocabulary_info = {
                    'vocabulary_id': vocabulary_id,
                    'vocabulary': self.vocabulary_cache[vocabulary_id],
                    'status': 'OK',
                    'prefix': self._extract_prefix_from_id(vocabulary_id)
                }
                print(f"从缓存获取热词表成功，ID: {vocabulary_id}")
                return vocabulary_info

            # 尝试通过query_vocabulary方法获取热词表内容
            try:
                query_response = self.service.query_vocabulary(vocabulary_id)

                # 处理query_vocabulary的返回结果
                if isinstance(query_response, dict):
                    vocabulary_data = query_response.get('vocabulary', [])

                # 构建返回信息
                vocabulary_info = {
                    'vocabulary_id': vocabulary_id,
                    'vocabulary': vocabulary_data,
                    'status': 'OK',
                    'prefix': self._extract_prefix_from_id(vocabulary_id)
                }

                # 更新缓存
                if isinstance(vocabulary_data, list):
                    self.vocabulary_cache[vocabulary_id] = vocabulary_data

                print(f"通过API获取热词表成功，ID: {vocabulary_id}")
                return vocabulary_info

            except Exception as query_error:
                print(f"query_vocabulary方法调用失败: {query_error}")

                # 如果query_vocabulary失败，返回基本信息
                vocabulary_info = {
                    'vocabulary_id': vocabulary_id,
                    'vocabulary': [],
                    'status': 'UNKNOWN',
                    'prefix': self._extract_prefix_from_id(vocabulary_id),
                    'error': str(query_error)
                }
                return vocabulary_info

        except Exception as e:
            print(f"获取热词表失败: {e}")
            return None

    def list_vocabularies(self, prefix: Optional[str] = None, page_index: int = 0, page_size: int = 10) -> List[
        Dict[str, Any]]:
        """
        查询所有热词表
        
        Args:
            prefix: 热词表前缀过滤，None表示获取所有
            page_index: 页码，从0开始
            page_size: 每页大小
            
        Returns:
            热词表列表
        """
        try:
            # 查询热词表列表
            params = {
                'page_index': page_index,
                'page_size': page_size
            }
            if prefix:
                params['prefix'] = prefix

            response = self.service.list_vocabularies(**params)

            # 处理返回的响应格式
            vocabulary_list = []

            # 检查是否是包含多个热词表的列表格式
            for item in response:
                if isinstance(item, dict) and 'vocabulary_id' in item:
                    vocab_info = {
                        'vocabulary_id': item.get('vocabulary_id'),
                        'gmt_create': item.get('gmt_create'),
                        'gmt_modified': item.get('gmt_modified'),
                        'status': item.get('status'),
                        'prefix': self._extract_prefix_from_id(item.get('vocabulary_id', ''))
                    }
                    vocabulary_list.append(vocab_info)

            print(f"查询到 {len(vocabulary_list)} 个热词表")
            return vocabulary_list

        except Exception as e:
            print(f"查询热词表列表失败: {e}")
            return []

    def _extract_prefix_from_id(self, vocabulary_id: str) -> str:
        """从热词表ID中提取前缀"""
        if not vocabulary_id:
            return ''

        # 热词表ID格式通常为: vocab-{prefix}-{hash}
        if vocabulary_id.startswith('vocab-'):
            parts = vocabulary_id.split('-')
            if len(parts) >= 2:
                return parts[1]

        return ''

    def add_words_to_vocabulary(self, vocabulary_id: str, new_words: List[dict]) -> bool:
        """
        向现有热词表添加新词汇
        
        Args:
            vocabulary_id: 热词表ID
            new_words: 要添加的新词汇列表
            
        Returns:
            是否添加成功
        """
        try:
            # 验证新词汇
            self._validate_vocabulary(new_words)

            # 获取现有热词表
            if vocabulary_id in self.vocabulary_cache:
                current_vocabulary = self.vocabulary_cache[vocabulary_id]
            else:
                vocab_info = self.get_vocabulary(vocabulary_id)
                if not vocab_info:
                    print(f"热词表 {vocabulary_id} 不存在")
                    return False
                current_vocabulary = vocab_info.get('vocabulary', [])

            # 合并词汇（去重）
            existing_texts = {word['text'] for word in current_vocabulary}
            for word in new_words:
                if word['text'] not in existing_texts:
                    current_vocabulary.append(word)
                    existing_texts.add(word['text'])

            # 更新热词表
            return self.update_vocabulary(vocabulary_id, current_vocabulary)

        except Exception as e:
            print(f"添加词汇失败: {e}")
            return False

    def remove_words_from_vocabulary(self, vocabulary_id: str, words_to_remove: List[str]) -> bool:
        """
        从热词表中删除指定词汇
        
        Args:
            vocabulary_id: 热词表ID
            words_to_remove: 要删除的词汇文本列表
            
        Returns:
            是否删除成功
        """
        try:
            # 获取现有热词表
            if vocabulary_id in self.vocabulary_cache:
                current_vocabulary = self.vocabulary_cache[vocabulary_id]
            else:
                vocab_info = self.get_vocabulary(vocabulary_id)
                if not vocab_info:
                    print(f"热词表 {vocabulary_id} 不存在")
                    return False
                current_vocabulary = vocab_info.get('vocabulary', [])

            # 移除指定词汇
            words_to_remove_set = set(words_to_remove)
            filtered_vocabulary = [
                word for word in current_vocabulary
                if word['text'] not in words_to_remove_set
            ]

            # 更新热词表
            return self.update_vocabulary(vocabulary_id, filtered_vocabulary)

        except Exception as e:
            print(f"删除词汇失败: {e}")
            return False

    def _validate_vocabulary(self, vocabulary: List[dict]) -> None:
        """
        验证热词表格式
        
        Args:
            vocabulary: 热词表字典列表
            
        Raises:
            ValueError: 如果格式不正确
        """
        if not isinstance(vocabulary, list):
            raise ValueError("热词表必须是列表格式")

        if not vocabulary:
            raise ValueError("热词表不能为空")

        for i, word in enumerate(vocabulary):
            if not isinstance(word, dict):
                raise ValueError(f"热词表第{i + 1}项必须是字典格式")

            if 'text' not in word:
                raise ValueError(f"热词表第{i + 1}项缺少 'text' 字段")

            if 'lang' not in word:
                raise ValueError(f"热词表第{i + 1}项缺少 'lang' 字段")

            # weight 是可选的，默认为1
            if 'weight' not in word:
                word['weight'] = 1

    def _validate_prefix(self, prefix: str) -> None:
        """
        验证热词表前缀
        
        Args:
            prefix: 热词表前缀
            
        Raises:
            ValueError: 如果格式不正确
        """
        if not isinstance(prefix, str):
            raise ValueError("前缀必须是字符串")

        if len(prefix) >= 10:
            raise ValueError("前缀长度必须小于10个字符")

        if not prefix.replace('_', '').isalnum() or not prefix.islower():
            raise ValueError("前缀只能包含小写字母、数字和下划线")

    def create_default_vocabulary(self, target_model: str = "paraformer-realtime-v2", prefix: str = "default") -> str:
        """
        创建默认热词表
        
        Args:
            target_model: 目标模型
            prefix: 前缀
            
        Returns:
            热词表ID
        """
        return self.create_vocabulary(target_model, prefix, MY_VOCABULARY)

    def initialize_vocabulary(self):
        """初始化热词表，使用VocabularyManager管理"""
        try:
            # 使用VocabularyManager创建热词表
            vocabulary_id = self.create_vocabulary("paraformer-realtime-v2", 'asr1', MY_VOCABULARY)
            logger.info(f"热词表初始化成功，ID: {vocabulary_id}")
            return vocabulary_id
        except Exception as e:
            logger.error(f"热词表初始化失败: {e}")
            # 如果热词表创建失败，返回None，识别引擎将不使用热词表
            return None


def main():
    """主函数 - 演示热词管理功能"""
    print("=== DashScope 热词管理器演示 ===\n")

    try:
        manager = VocabularyManager()

        # 1. 查询现有热词表
        print("1. 查询现有热词表...")
        vocabularies = manager.list_vocabularies(prefix="asr1", page_index=0, page_size=10)

        if vocabularies:
            print("\n=== 当前热词表列表 ===")
            for i, vocab in enumerate(vocabularies, 1):
                vocab_id = vocab.get('vocabulary_id', 'N/A')
                print(f"{i}. 热词表信息:")
                print(f"   ID: {vocab_id}")
        else:
            print("未找到任何热词表")

        # 2. 获取特定热词表详情
        if vocabularies and len(vocabularies) > 0:
            first_vocab_id = vocabularies[0]['vocabulary_id']
            print(f"2. 获取热词表详情 (ID: {first_vocab_id})...")
            vocab_detail = manager.get_vocabulary(first_vocab_id)

            if vocab_detail:
                print("热词表详情:")
                print(f"  ID: {vocab_detail.get('vocabulary_id')}")
                print(f"  前缀: {vocab_detail.get('prefix')}")
                print(f"  状态: {vocab_detail.get('status')}")

                vocabulary_data = vocab_detail.get('vocabulary', [])
                if isinstance(vocabulary_data, list) and vocabulary_data:
                    print(f"  热词数量: {len(vocabulary_data)}")
                    print("  热词列表:")
                    for word in vocabulary_data:
                        if isinstance(word, dict):
                            text = word.get('text', 'N/A')
                            weight = word.get('weight', 'N/A')
                            lang = word.get('lang', 'N/A')
                            print(f"    - {text} (权重: {weight}, 语言: {lang})")
                else:
                    print("  热词列表: 空或无法获取")

                # 如果有错误信息，显示出来
                if 'error' in vocab_detail:
                    print(f"  错误信息: {vocab_detail['error']}")
            else:
                print("获取热词表详情失败")

        # 删除热词表:
        # manager.delete_vocabulary('vocab-asr1-c8d82369dd714eba877b4ff10df5aa72')

        # 添加新词汇:
        # manager.add_words_to_vocabulary('vocab-asr1-c8d82369dd714eba877b4ff10df5aa72', [
        #     {'text': '新词', 'weight': 4, 'lang': 'zh'}
        # ])

    except Exception as e:
        logger.error(f"演示过程中发生错误: {e}")
        print(f"错误: {e}")

    print("\n=== 结束 ===")


if __name__ == "__main__":
    main()
